import re
import asyncio
import subprocess

from metagpt.actions import Action
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.logs import logger


class SimpleWriteCode(Action):
    # 在run方法中，我们需要声明当采取这个行动时，我们要对传入的内容做什么样的处理，
    # 在 SimpleWriteCode 类中，我们应该传入：“请你帮我写一个XXX的代码” 这样的字符串，也就是用户的输入，run方法需要对它进行处理，把他交给llm，等到llm返回生成结果后，我们再取出其中的代码部分返回。
    PROMPT_TEMPLATE: str = """
    Write a python function that can {instruction} and provide two runnnable test cases.
    Return ```python your_code_here ``` with NO other texts,
    your code:
    """

    name: str = "SimpleWriteCode"

    # 我们定义了一个 SimpleWriteCode 类，它继承自 Action类，我们重写了run 方法，该方法决定了我们对传入的内容到底要做什么样的处理。
    async def run(self, instruction: str):
        prompt = self.PROMPT_TEMPLATE.format(instruction=instruction)
        rsp = await self._aask(prompt)
        code_text = SimpleWriteCode.parse_code(rsp)
        return code_text

    @staticmethod
    def parse_code(rsp):
        pattern = r'```python(.*)```'
        match = re.search(pattern, rsp, re.DOTALL)
        code_text = match.group(1) if match else rsp
        return code_text


class SimpleCoder(Role):
    # 实现一个最简单的Role，只需要指定 Role 的 name（昵称） 和 profile（人设），并重写Role基类的 __init__ 与 _act 方法。
    name: str = "bondee"
    profile: str = "SimpleCoder"

    # __init__ 方法用来初始化这个Action，而 _act 方法决定了当这个角色行动时它的具体行动逻辑
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([SimpleWriteCode])

    # 在_act方法中，我们就会要求我们的智能体来执行这个动作，也就是我们需要调用todo.run()方法
    async def _act(self) -> Message:
        logger.info(f"{self._setting}: ready to {self.rc.todo}")
        todo = self.rc.todo  # todo will be SimpleWriteCode()

        msg = self.get_memories(k=1)[0]  # find the most recent messages

        code_text = await todo.run(msg.content)
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg

#
# async def main():
#     msg = "write a function that calculates the sum of a list"
#     role = SimpleCoder()
#     logger.info(msg)
#     result = await role.run(msg)
#     logger.info(result)
#
#
# asyncio.run(main())



